import torch
import numpy as np
import sys
sys.path.append("..") # 为了导入上层目录的d2lzh_pytorch
from basic_knowledge import d2lzh_pytorch as d2l

# 批量数据集大小为256
batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)

# 初始化参数
# 每个样本输入是高和宽均为28像素的图像,每个图像有10个类别
num_inputs = 784
num_outputs = 10
W = torch.tensor(np.random.normal(0, 0.01, (num_inputs, num_outputs)), dtype=torch.float)
b = torch.zeros(num_outputs, dtype=torch.float)
# 模型参数梯度
W.requires_grad_(requires_grad=True)
b.requires_grad_(requires_grad=True)
# softmax运算会先通过exp函数对每个元素做指数运算，再对exp矩阵同行元素求和
# 最后令矩阵每行各元素与该行元素之和相除，这样一来，最终得到的矩阵每行元素和为1且非负
# softmax运算的输出矩阵中的任意一行元素代表了一个样本在各个输出类别上的预测概率。
def softmax(X):
    X_exp = X.exp()
    partition = X_exp.sum(dim=1, keepdim=True)
    return X_exp / partition  # 这里应用了广播机制
# 定义模型
# 通过view函数将每张原始图像改成长度为num_inputs的向量
# 将二维的变成一维的，-1表示可以推测出来的参数
# 1*784矩阵乘784*10矩阵+1*10矩阵，输出一个图像的各个类别的概率
def net(X):
    return softmax(torch.mm(X.view((-1, num_inputs)), W) + b)
# 定义交叉熵损失函数，y是实际的标签，每个样本对应一个标量，代表类别号
# y_hat是预测结果，对应每个样本一个1*10的张量
def cross_entropy(y_hat, y):
    return -torch.log(y_hat.gather(1, y.view(-1, 1)))
# 计算分类准确率，分类准确率即正确预测数量与总预测数量之比
# y_hat.argmax(dim=1)返回矩阵y_hat每行中最大元素的索引
def accuracy(y_hat, y):
    return (y_hat.argmax(dim=1) == y).float().mean().item()

# 训练模型
num_epochs, lr = 5, 0.1
d2l.train_ch3(net, train_iter, test_iter, cross_entropy, num_epochs, batch_size, [W, b], lr)